#include "pch.h"
#include "msghook.h"
#include <Ws2tcpip.h>
#include <winsock2.h>
#include <nlohmann/json.hpp>
#include "common.h"
#include "offset.h"

using namespace nlohmann;
using namespace std;

#define CLIENT_IP "127.0.0.1"
static int kServerPort = 0;
static int kMessageHooked = FALSE;
static char kServerIp[16]= "127.0.0.1";
static DWORD kWeChatWinBase = GetWeChatWinBase();

static char kOriginSnsMsgAsmCode[5] = {0};
static char kOriginReceMsgAsmCode[5] = {0};
static DWORD kReceMsgJmpBackAddress = kWeChatWinBase + WX_RECV_MSG_HOOK_OFFSET + 0x5;
static DWORD kReceMsgNextAddress = kWeChatWinBase + WX_RECV_MSG_HOOK_NEXT_OFFSET;
static DWORD kSnsMsgJmpBackAddress = kWeChatWinBase + WX_SNS_HOOK_OFFSET + 0x5;
static DWORD kSnsMsgNextAddress = kWeChatWinBase + WX_SNS_HOOK_NEXT_OFFSET;

struct InnerMessageStruct {
  char *buffer;
  int length;
  ~InnerMessageStruct() {
    if (this->buffer != NULL) {
      delete[] this->buffer;
      this->buffer = NULL;
    }
  }
};

int CheckLogin() {
  kWeChatWinBase = GetWeChatWinBase();
  DWORD accout_service_addr = kWeChatWinBase + WX_ACCOUNT_SERVICE_OFFSET;
  DWORD service_addr = NULL;
  __asm {
      PUSHAD 
      CALL       accout_service_addr                          
      MOV        service_addr,EAX
      POPAD
  }
  if (service_addr) {
    return *(DWORD *)(service_addr + 0x4C8);
  }
  else {
    return 0;
  }
}

/// @brief send message  by socket
/// @param buffer content
/// @param len len
/// @return true/false
BOOL SendBySocket(const char* buffer, size_t len) {
    if (kServerPort == 0) {
        return false;
    }
    SOCKET client_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (client_socket < 0) {
#ifdef _DEBUG
        cout << "create socket error,"
            << " errno:" << errno << endl;
#endif
        return false;
    }
    BOOL status = false;
    sockaddr_in client_addr;
    memset(&client_addr, 0, sizeof(client_addr));
    client_addr.sin_family = AF_INET;
    client_addr.sin_port = htons((u_short)kServerPort);
    InetPtonA(AF_INET, kServerIp, &client_addr.sin_addr.s_addr);
    if (connect(client_socket, reinterpret_cast<sockaddr*>(&client_addr),
        sizeof(sockaddr)) < 0) {
#ifdef _DEBUG
        cout << "kServerIp:" << kServerIp << endl;
        cout << "port:" << kServerPort << endl;
        cout << "connect error,"
            << " errno:" << errno << endl;
#endif
        return false;
    }
    char recv_buf[1024] = { 0 };
    int ret = send(client_socket, buffer, len, 0);
    if (ret == -1 || ret == 0) {
#ifdef _DEBUG
        cout << "send fail,"
            << " errno:" << errno << endl;
#endif
        closesocket(client_socket);
        return false;
    }
    memset(recv_buf, 0, sizeof(recv_buf));
    ret = recv(client_socket, recv_buf, sizeof(recv_buf), 0);
    closesocket(client_socket);
    if (ret == -1 || ret == 0) {
#ifdef _DEBUG
        cout << "the server close" << endl;
#endif
        return false;
    }
    return true;
}

/// @brief send wrap
/// @param msg msg
void SendSocketMessage(InnerMessageStruct* msg) {
    if (msg == NULL) {
        return;
    }
    unique_ptr<InnerMessageStruct> sms(msg);
    json j_msg =
        json::parse(msg->buffer, msg->buffer + msg->length, nullptr, false);
    if (j_msg.is_discarded() == true) {
        return;
    }
    string jstr = j_msg.dump() + "\n";
#ifdef _DEBUG
    cout << "json:" + jstr << endl;
#endif
    SendBySocket(jstr.c_str(), jstr.size());
}

/// @brief msg handle
/// @param msg_addr msg address in memory
void __cdecl OnRecvMsg(DWORD msg_addr) {
  json j_msg;
  unsigned long long msgid = *(unsigned long long *)(msg_addr + 0x30);
  j_msg["msgId"] = msgid;
  j_msg["pid"] = GetCurrentProcessId();
  j_msg["type"] = *(DWORD *)(msg_addr + 0x38);
  j_msg["isSendMsg"] = *(BOOL *)(msg_addr + 0x3C);
  if (j_msg["isSendMsg"].get<BOOL>()) {
    j_msg["isSendByPhone"] = (int)(*(BYTE *)(msg_addr + 0xD8));
  }
  j_msg["time"] =
      unicode_to_utf8((wchar_t *)GetTimeW(*(DWORD *)(msg_addr + 0x44)).c_str());
  j_msg["timestamp"] = *(DWORD *)(msg_addr + 0x44);
  j_msg["fromGroup"] =
      unicode_to_utf8((wchar_t *)READ_WSTRING(msg_addr, 0x48).c_str());
  int length = *(DWORD *)(msg_addr + 0x178);
  if (length == 0) {
    j_msg["fromUser"] = j_msg["fromGroup"].get<std::string>();
  } else {
    j_msg["fromUser"] =
        unicode_to_utf8((wchar_t *)READ_WSTRING(msg_addr, 0x174).c_str());
  }
  int content_len = *(DWORD *)(msg_addr + 0x74);
  if (content_len > 0) {
    j_msg["content"] =
        unicode_to_utf8((wchar_t *)READ_WSTRING(msg_addr, 0x70).c_str());
  }

  string jstr = j_msg.dump() + '\n';
  InnerMessageStruct *inner_msg = new InnerMessageStruct;
  inner_msg->buffer = new char[jstr.size() + 1];
  memcpy(inner_msg->buffer, jstr.c_str(), jstr.size() + 1);
  inner_msg->length = jstr.size();
  HANDLE thread = CreateThread(
      NULL, 0, (LPTHREAD_START_ROUTINE)SendSocketMessage, inner_msg, NULL, 0);
  if (thread) {
    CloseHandle(thread);
  }
}

 void __cdecl OnSnsTimeLineMsg(DWORD msg_addr) {
  json j_sns;
  DWORD begin_addr = *(DWORD *)(msg_addr + 0x20);
  DWORD end_addr = *(DWORD *)(msg_addr + 0x24);
  #ifdef _DEBUG
   cout << "begin" <<begin_addr<<endl;
   cout << "end" <<end_addr<<endl;
  #endif
  if (begin_addr == 0) {
    j_sns = {{"data", json::array()}};
  } else {
    while (begin_addr < end_addr) {
      json j_item;
      j_item["snsId"] = *(unsigned long long *)(begin_addr);
      j_item["createTime"] = *(DWORD *)(begin_addr + 0x2C);

      j_item["senderId"] =
          unicode_to_utf8((wchar_t *)READ_WSTRING(begin_addr, 0x18).c_str());

      j_item["content"] =
          unicode_to_utf8((wchar_t *)READ_WSTRING(begin_addr, 0x3c).c_str());

      j_item["xml"] =
          unicode_to_utf8((wchar_t *)READ_WSTRING(begin_addr, 0x384).c_str());

      j_sns["data"].push_back(j_item);
      begin_addr += 0xB48;
    }
  }
  string jstr = j_sns.dump() + '\n';
  InnerMessageStruct *inner_msg = new InnerMessageStruct;
  inner_msg->buffer = new char[jstr.size() + 1];
  memcpy(inner_msg->buffer, jstr.c_str(), jstr.size() + 1);
  inner_msg->length = jstr.size();
  HANDLE thread = CreateThread(
      NULL, 0, (LPTHREAD_START_ROUTINE)SendSocketMessage, inner_msg, NULL, 0);
  if (thread) {
    CloseHandle(thread);
  }
}


/// @brief  hook sns msg implement
_declspec(naked) void handle_sns_msg() {
  __asm {
    PUSHAD
		PUSHFD
    PUSH    [ESP + 0x24]
		CALL    OnSnsTimeLineMsg
		ADD     ESP, 0x4
    POPFD
		POPAD
		CALL    kSnsMsgNextAddress
		JMP     kSnsMsgJmpBackAddress
  }
}

/// @brief  hook implement
_declspec(naked) void handle_sync_msg() {
  __asm {
    PUSHAD
		PUSHFD
		PUSH    ECX
		CALL    OnRecvMsg
		ADD     ESP, 0x4
		POPFD
		POPAD
		CALL    kReceMsgNextAddress
		JMP     kReceMsgJmpBackAddress
  }
}

/// @brief hook any address   address+0x5
/// @param port 端口
/// @return 成功返回1,失败返回0
int HookRecvMsg(char* client_ip,int port) {
  kServerPort = port;
  strcpy_s(kServerIp,client_ip);
  kWeChatWinBase = GetWeChatWinBase();

  DWORD hook_recv_msg_addr = kWeChatWinBase + WX_RECV_MSG_HOOK_OFFSET;
  kReceMsgNextAddress = kWeChatWinBase + WX_RECV_MSG_HOOK_NEXT_OFFSET;
  kReceMsgJmpBackAddress = hook_recv_msg_addr + 0x5;
  HookAnyAddress(hook_recv_msg_addr, (LPVOID)handle_sync_msg,
                 kOriginReceMsgAsmCode);

  DWORD hook_sns_msg_addr = kWeChatWinBase + WX_SNS_HOOK_OFFSET;
  kSnsMsgNextAddress = kWeChatWinBase + WX_SNS_HOOK_NEXT_OFFSET;
  kSnsMsgJmpBackAddress = hook_sns_msg_addr + 0x5;
  HookAnyAddress(hook_sns_msg_addr, (LPVOID)handle_sns_msg,
                 kOriginSnsMsgAsmCode);
  

  kMessageHooked = TRUE;
  return 1;
}
